/* vim: ft=c */
#ifndef BALLOC_INCLUDED
#define BALLOC_INCLUDED
/******************************************************************************
* (c)2025 Michael T. Richter
*
* This software is distributed under the terms of WTFPLv2.  The full terms and
* text of the license can be found at http://www.wtfpl.net/txt/copying
******************************************************************************/
/** @file
 *  @brief A multi-arena bump allocator system.
 *
 * A frequent pattern in dynamic memory usage is allocation in dribs and drabs
 * sequentially followed by a mass freeing.  Using traditional `malloc()` and
 * related functions for this is wasteful of MCU resources, wasteful of memory
 * for all the bookkeeping functionality involved, and tediously error-prone in
 * subsequent cleanup.
 *
 * This library provides an implementation of a multi-arena bump allocator
 * scheme that groups sequential allocations in a single block of memory and
 * frees all of them in a single operation.
 *
 * Important philosophical note: this library is built on a "don't supply what
 * you probably don't need" basis.  It is also distributed under WTFPLv2 which
 * means you can do to it whatever you like and I don't care.  The specifics of
 * what I don't care about here include:
 *
 *  1. Pre-emption.  I only do work with main loops, reactors, and coroutines.
 *     I don't do pre-emptive multi-tasking because it's dangerous.  As such I
 *     have not littered my code with hooks for mutexes, critical sections, spin
 *     locks and other such pre-emptive shitfuckery.  If you need these, use the
 *     power of "What The Fuck" from WTFPLv2 and put 'em in.  It's your code now
 *     do with it what you will.
 *  2. Compiler magic.  There's no `restrict` or `__attribute__((malloc))` or
 *     other such magical incantations here.  My target environment is embedded
 *     systems and, let's be honest, these incantations are not going to make any
 *     kind of noticeable difference in a sane embedded environment.  Certainly
 *     not sufficiently noticeable for me to want to start that dangerous dark
 *     path of macromania where I put BALLOC_ATTRIBUTE_MALLOC in front of half
 *     the lines of code that will expand to nothing in most cases anyway.  If
 *     you think these compiler magic details are important, use the power of
 *     "What The Fuck" and add them.
 *  3. Unnecessary optimizations.  I'm making this "fast enough" for my
 *     purposes.  If you need faster, you know the solution by now.
 *  4. Whining about "professionalism".  The project is licensed literally under
 *     the "What The Fuck Public Licence v2".  I think professionalism has long
 *     been thrown out the window from the 52nd storey.  You're not paying me
 *     for this, so by definition this is not a professional product.  If you
 *     want "professionalism" in my documentation and comments, hire me.  Or
 *     don't.  I'm easy either way.
 *
 *  The ultimate answer to any problem you might have is:
 *
 *                      THE POWER OF WTFPLV2 ALLOWS YOU!
 */

#include <stdbool.h>
#include <stddef.h>

/** @defgroup balloc Bump allocation system.
 *
 * @{
 */

#ifndef UNUSED
#define UNUSED(X) ((void)(X))
#endif

#ifndef BALLOC_ALIGNMENT
#define BALLOC_ALIGNMENT (sizeof(void*))
#endif
typedef char BALLOC_ALIGNMENT_check[(BALLOC_ALIGNMENT & (BALLOC_ALIGNMENT - 1)) == 0 ? 1 : -1];
// If the typedef fails to compile, you're doing something weird.

#ifndef BALLOC_SECURE
#define BALLOC_SECURE false
#endif

/** @brief Opaque allocator handle permitting multiple object pools. */
typedef struct _balloc_arena *balloc_arena;

/** @brief Stats for measuring and debugging allocators. */
typedef struct balloc_arena_stats
{
    size_t current_size;    /**< the current allocated size of the arena  */
    size_t current_used;    /**< the allocated size actually in use       */
    int    error_count;     /**< the count of failed allocations          */
} balloc_arena_stats;

/** @brief Error callback for allocator errors.
 *
 * @param [in] file          The source file where the error was discovered.
 * @param [in] line          The line of the file where the error was discovered.
 * @param [in] error_message The error's textual message.
 *
 * A function pointer
 */
typedef void (*balloc_error_handler_t)(const char *file, const int line, const char *error_message);
extern balloc_error_handler_t balloc_error_handler;

/** @brief Create a new balloc arena.
 *
 * @param [in] size The number of bytes to allocate for the pool.
 * @param [in] pool For cases where a statically-allocated pool is desired, a pointer to it.
 *
 * @return The arena handle used in subsequent API calls, NULL on failure.
 *
 * @note The `balloc_error_handler()` callback is called before NULL is returned.
 */
balloc_arena balloc_new_arena(size_t size);
balloc_arena balloc_new_static_arena(void *pool, size_t size);

/** @brief Allocate memory from the given arena.
 *
 * @param [in,out] arena The handle of the arena to allocate from; will update so not const.
 *
 * @return A pointer to the memory space, NULL on failure.
 *
 * @note The `balloc_error_handler()` callback is called before NULL is returned.
 */
void *balloc(balloc_arena arena, size_t bytes);

/** @brief Clear **ALL** memory associated with the given arena and reset it to
 *         a pristine state.
 *
 * @param [in,out] arena The handle of the arena to clear.  Allocated
 *                       memory is zeroed.  Arena data is also cleared.
 */
void balloc_reset(balloc_arena arena);

/** @brief Clear and free **ALL** memory associated with the given arena.
 *
 * @param [in,out] arena The handle of the arena to clear and free.  Allocated
 *                       memory is zeroed, then freed.  Arena data is then
 *                       also cleared.
 */
void balloc_free(balloc_arena arena);

/** @brief Collect some allocator stats for measurement and debugging.
 *
 * @param [in] allocator The allocator to collect stats from.
 *
 * @return A structure with some allocator statistics for measurement and
 *         debugging.
 */
balloc_arena_stats balloc_get_arena_stats(const balloc_arena arena);

#ifdef BALLOC_IMPLEMENTATION
#include <stdint.h>
#include <stdlib.h>
#include <string.h>

/*******************************************************************************
* Static helper function declarations
*******************************************************************************/


/*******************************************************************************
* Opaque data structures
*******************************************************************************/

struct _balloc_arena
{
    bool    cleanup;      /**< the pool needs to be cleaned up         */
    size_t  current_size; /**< the current allocated size of the arena */
    size_t  current_used; /**< the allocated size actually in use      */
    int     error_count;  /**< the count of failed allocations         */
    uint8_t *pool;        /**< the actual memory pool of this arena    */
};

/*******************************************************************************
* API function definitions
*******************************************************************************/
balloc_arena balloc_new_arena(size_t size)
{
    struct _balloc_arena *rv = NULL;
    uint8_t *pool = malloc(size);

    if (pool)
    {
        rv = malloc(sizeof(struct _balloc_arena));

        if (rv)
        {
            rv->current_size = size;
            rv->current_used = rv->error_count = 0;
            rv->pool = pool;
            rv->cleanup = true;
        }
        else
        {
            free(pool);
        }
    }

    return rv;
}

balloc_arena balloc_new_static_arena(void *pool, size_t size)
{
    struct _balloc_arena *rv = NULL;

    if (pool)
    {
        rv = malloc(sizeof(struct _balloc_arena));

        if (rv)
        {
            rv->current_size = size;
            rv->current_used = rv->error_count = 0;
            rv->pool = pool;
            rv->cleanup = false;
        }
    }

    return rv;
}

void *balloc(balloc_arena arena, size_t bytes)
{
    if (!arena) { return NULL; }

    uint8_t *rv;
    size_t abytes = (bytes + (BALLOC_ALIGNMENT - 1)) & ~(BALLOC_ALIGNMENT - 1);

    if ((arena->current_size - arena->current_used) >= abytes)
    {
        // we have enough bytes, so return the current position and bump it
        rv = arena->pool + arena->current_used;
        arena->current_used += abytes;
    }
    else
    {
        // we can't allocate, so signal the error and return
        // NULL; the arena is still usable for smaller allocation
        balloc_error_handler(__FILE__, __LINE__, "arena not large enough for desired size");
        arena->error_count++;
        rv = NULL;
    }

    return rv;
}

void balloc_reset(balloc_arena arena)
{
    if (arena)
    {
        if (BALLOC_SECURE) { memset(arena->pool, 0, arena->current_size); }
        arena->current_used = 0;
        arena->error_count = 0;
    }
}

void balloc_free(balloc_arena arena)
{
    if (arena)
    {
        if (BALLOC_SECURE)  { memset(arena->pool, 0, arena->current_size); }
        if (arena->cleanup) { free(arena->pool); }
        memset(arena, 0, sizeof(struct _balloc_arena));
        free(arena);
    }
}

balloc_arena_stats balloc_get_arena_stats(const balloc_arena arena)
{
    balloc_arena_stats rv = { 0 };

    if (arena)
    {
        rv.current_size = arena->current_size;
        rv.current_used = arena->current_used;
        rv.error_count = arena->error_count;
    }

    return rv;
}

/*******************************************************************************
* Static helper function definitions
*******************************************************************************/

void default_balloc_error_handler(const char *file, const int line, const char *error_message)
{
    UNUSED(file);
    UNUSED(line);
    UNUSED(error_message);
}
balloc_error_handler_t balloc_error_handler = default_balloc_error_handler;

#endif //BALLOC_IMPLEMENTATION

#endif //!BALLOC_INCLUDED
